from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import SubprocVecEnv, VecFrameStack
from stable_baselines3.common.callbacks import EvalCallback
from obs_test import env_make


def main():

    vec_env = SubprocVecEnv([env_make for _ in range(4)])
    vec_env = VecFrameStack(vec_env, 4, channels_order='last')  # 帧叠加
    eval_callback = EvalCallback(vec_env, best_model_save_path='./best_model', log_path='./callback_logs', eval_freq=10000//4)

    model = PPO('CnnPolicy', vec_env, verbose=1, tensorboard_log='logs',
                learning_rate=1e-4,
                n_steps=2048,
                batch_size=512,
                n_epochs=10,
                ent_coef=0.1,
                target_kl=0.1,
                )

    model.learn(total_timesteps=1e6, callback=eval_callback)

    model.save('ppo_mario')


if __name__ == '__main__':
    main()
